using BepuUtilities;
using BepuUtilities.Collections;
using BepuUtilities.Memory;
using BepuPhysics.Constraints;
using System;
using System.Diagnostics;
#if MYCODE
using BepuUtilities.Vectors;
#else
using System.Numerics;
#endif
namespace BepuPhysics
{
    public class ConstraintLayoutOptimizer
    {
        Bodies bodies;
        Solver solver;
        struct Optimization
        {
            /// <summary>
            /// 要优化的目标约束束的索引。
            /// </summary>
            public int BundleIndex;
            /// <summary>
            /// 上次优化类型批处理的索引。
            /// </summary>
            public int TypeBatchIndex;
            /// <summary>
            /// 上次优化的批处理的索引。
            /// </summary>
            public int BatchIndex;

        }

        Optimization nextTarget;

        /// <summary>
        /// 如果为True,区域将偏移半个区域宽度。已切换每一帧。偏移量允许排序的区域混合,最终收敛为完全排序。
        /// </summary>
        bool shouldOffset;

        float optimizationFraction;
        public float OptimizationFraction
        {
            get { return optimizationFraction; }
            set
            {
                if (value < 0 || value > 1)
                    throw new ArgumentException("Optimization fraction must be from 0 to 1.");
                optimizationFraction = value;
            }
        }

        Action<int> generateSortKeysDelegate;
        Action<int> regatherDelegate;
        Action<int> copyToCacheAndSortDelegate;

        public ConstraintLayoutOptimizer(Bodies bodies, Solver solver, float optimizationFraction = 0.044f)
        {
            this.bodies = bodies;
            this.solver = solver;
            OptimizationFraction = optimizationFraction;

            generateSortKeysDelegate = GenerateSortKeys;
            regatherDelegate = Regather;
            copyToCacheAndSortDelegate = CopyToCacheAndSort;
        }

        void Wrap(ref Optimization o)
        {
            ref var activeSet = ref solver.ActiveSet;
            Debug.Assert(activeSet.Batches.Count > 0, "Shouldn't be trying to optimize zero constraints.");
            while (true)
            {
                if (o.BatchIndex >= activeSet.Batches.Count)
                {
                    o = new Optimization();
                }
                else if (o.TypeBatchIndex >= activeSet.Batches[o.BatchIndex].TypeBatches.Count)
                {
                    // 最后一个约束批次之前的批次可能缺少约束。那样的话,试试下一个吧。
                    ++o.BatchIndex;
                    o.TypeBatchIndex = 0;
                    o.BundleIndex = 0;
                }
                else if (o.BundleIndex >= activeSet.Batches[o.BatchIndex].TypeBatches[o.TypeBatchIndex].BundleCount)
                {
                    ++o.TypeBatchIndex;
                    o.BundleIndex = 0;
                }
                else
                {
                    break;
                }
            }
        }

        Optimization FindOffsetFrameStart(Optimization o, int maximumRegionSizeInBundles)
        {
            Wrap(ref o);

            ref var activeSet = ref solver.ActiveSet;
            var spaceRemaining = activeSet.Batches[o.BatchIndex].TypeBatches[o.TypeBatchIndex].BundleCount - o.BundleIndex;
            if (spaceRemaining <= maximumRegionSizeInBundles)
            {
                ++o.TypeBatchIndex;
                Wrap(ref o);
            }
            // 请注意,捆绑包计数不会被缓存;以上类型的批处理可能会有所不同。
            o.BundleIndex = Math.Max(0,
                Math.Min(
                    o.BundleIndex + maximumRegionSizeInBundles / 2,
                    activeSet.Batches[o.BatchIndex].TypeBatches[o.TypeBatchIndex].BundleCount - maximumRegionSizeInBundles));

            return o;
        }

        public unsafe void Update(BufferPool bufferPool, IThreadDispatcher threadDispatcher = null)
        {
            // TODO：与设置多线程相关的成本可能超过较小模拟的实际优化成本。
            // 根据一些经验测试,您可能希望退回到单线程。
            // 如果没有约束,优化就没有意义-这是一个必要的测试,因为我们稍后假设0是一个有效的批处理索引。
            ref var activeSet = ref solver.ActiveSet;
            if (activeSet.Batches.Count == 0)
                return;
            var regionSizeInBundles = (int)Math.Max(2, Math.Round(activeSet.BundleCount * optimizationFraction));
            // 捆绑中的区域大小应该可以被2整除,以便可以将其偏移一半。
            if ((regionSizeInBundles & 1) == 1)
                ++regionSizeInBundles;
            // 请注意,我们要求所有区域都是捆绑对齐的。这对于TypeBatch排序过程很重要,该过程倾向于使用从捆绑包数组到缓存的批量复制。
            // 如果不捆绑对齐,这些批量复制将由于约束AOSOA布局而变得复杂。

            Optimization target;
            if (shouldOffset)
            {
                // 使用上一帧的开始创建新目标。
                target = FindOffsetFrameStart(nextTarget, regionSizeInBundles);
                Debug.Assert(activeSet.Batches[target.BatchIndex].TypeBatches[target.TypeBatchIndex].BundleCount <= regionSizeInBundles || target.BundleIndex != 0,
                    "On offset frames, the only time a target bundle can be 0 is if the batch is too small for it to be anything else.");
                // Console.WriteLine($"针对{target.BatchIndex}.{target.TypeBatchIndex}：{target.BundleIndex}");的偏移帧
                Debug.Assert(activeSet.Batches[target.BatchIndex].TypeBatches.Count > target.TypeBatchIndex);
            }
            else
            {
                // 由于约束集自上次执行以来可能已任意更改,因此请从批处理向下验证。
                target = nextTarget;
                Wrap(ref target);
                Debug.Assert(activeSet.Batches[target.BatchIndex].TypeBatches.Count > target.TypeBatchIndex);
                nextTarget = target;
                nextTarget.BundleIndex += regionSizeInBundles;
                // Console.WriteLine($"针对{target.BatchIndex}.{target.TypeBatchIndex}：{target.BundleIndex}");的普通帧
            }
            // 请注意,我们在多个帧上有两个独立的并行优化。在每个帧的基础上在它们之间交替是一种相当简单的方式来保证
            // 最终在排序中收敛。我们只推进非偏移版本;偏移位置基于非偏移版本的最后开始位置。
            shouldOffset = !shouldOffset;


            var maximumRegionSizeInConstraints = regionSizeInBundles * Vector<int>.Count;

            var typeBatch = activeSet.Batches[target.BatchIndex].TypeBatches.GetPointer(target.TypeBatchIndex);
            SortByBodyLocation(typeBatch, target.BundleIndex, Math.Min(typeBatch->ConstraintCount - target.BundleIndex * Vector<int>.Count, maximumRegionSizeInConstraints),
                solver.HandleToConstraint, bodies.ActiveSet.Count, bufferPool, threadDispatcher);

        }

        unsafe TypeBatch* typeBatchPointer;
        IThreadDispatcher threadDispatcher;
        Buffer<ConstraintLocation> handlesToConstraints;
        struct MultithreadingContext
        {
            public Buffer<int> SortKeys;
            public Buffer<int> SourceIndices;
            public RawBuffer BodyReferencesCache;
            public int SourceStartBundleIndex;
            public int BundlesPerWorker;
            public int BundlesPerWorkerRemainder;
            public int TypeBatchConstraintCount;

            public Buffer<int> SortedKeys; // 这实际上只是存储以供调试使用。
            public Buffer<int> SortedSourceIndices;
            public Buffer<int> ScratchKeys;
            public Buffer<int> ScratchValues;
            public Buffer<ConstraintHandle> IndexToHandleCache;
            public RawBuffer PrestepDataCache;
            public RawBuffer AccumulatesImpulsesCache;
            public int KeyUpperBound;
            public int ConstraintsInSortRegionCount;
            // 请注意,这些不同于阶段1-排序中的一个线程专用于排序。这些都是关于剩余的线索。
            public int CopyBundlesPerWorker;
            public int CopyBundlesPerWorkerRemainder;
        }
        MultithreadingContext context;

        unsafe void GenerateSortKeys(int workerIndex)
        {
            var localWorkerBundleStart = context.BundlesPerWorker * workerIndex + Math.Min(workerIndex, context.BundlesPerWorkerRemainder);
            var workerBundleStart = context.SourceStartBundleIndex + localWorkerBundleStart;
            var workerBundleCount = workerIndex < context.BundlesPerWorkerRemainder ? context.BundlesPerWorker + 1 : context.BundlesPerWorker;
            var workerConstraintStart = workerBundleStart << BundleIndexing.VectorShift;
            // 请注意,我们可以迭代的约束数量受类型Batch的约束计数的限制。最后一个捆绑包可能未满。
            var workerConstraintCount = Math.Min(context.TypeBatchConstraintCount - workerConstraintStart, workerBundleCount << BundleIndexing.VectorShift);
            if (workerConstraintCount <= 0)
                return; // 没有剩余的工作。
            var localWorkerConstraintStart = localWorkerBundleStart << BundleIndexing.VectorShift;

            ref var typeBatch = ref *typeBatchPointer;
            solver.TypeProcessors[typeBatch.TypeId].GenerateSortKeysAndCopyReferences(ref typeBatch,
                workerBundleStart, localWorkerBundleStart, workerBundleCount,
                workerConstraintStart, localWorkerConstraintStart, workerConstraintCount,
                ref context.SortKeys[localWorkerConstraintStart], ref context.SourceIndices[localWorkerConstraintStart], ref context.BodyReferencesCache);
        }

        unsafe void CopyToCacheAndSort(int workerIndex)
        {
            // 排序只需要准备好排序关键字和索引。缓存还不需要执行。
            // 考虑到排序已经非常快,并且尝试独立地多线程不是一个好主意,我们将把它捆绑在一起
            // 剩余的缓存副本。此阶段非常受内存限制,排序可能与复制持续时间不匹配,但没有
            // 在这些时间尺度(<150us)下可以容纳复杂的方案。我们只是试着用几个简单的把戏尽可能地获得更多的利益。
            // 在带宽/计算比类似于3770K和1600 MHz内存的计算机上,我们最有可能获得的加速不会超过2.5倍。
            if (workerIndex == threadDispatcher.ThreadCount - 1)
            {
                // TODO：如果这里最终成为您实际使用线程内存池的唯一位置,那么您最好还是去掉它
                // 支持仅预先分配每个1024个整数的workerCount缓冲区。它最初用于创建特定于TypeBatch的内存的用法已不复存在。
                LSBRadixSort.Sort(
                    ref context.SortKeys, ref context.SourceIndices,
                    ref context.ScratchKeys, ref context.ScratchValues, 0, context.ConstraintsInSortRegionCount,
                    context.KeyUpperBound, threadDispatcher.GetThreadMemoryPool(workerIndex),
                    out context.SortedKeys, out context.SortedSourceIndices);
            }
            // 请注意,如果池中只有一个线程,则辅助进程0仍会进行复制。主要用于调试目的。
            if (threadDispatcher.ThreadCount == 1 || workerIndex < threadDispatcher.ThreadCount - 1)
            {
                var localWorkerBundleStart = context.CopyBundlesPerWorker * workerIndex + Math.Min(workerIndex, context.CopyBundlesPerWorkerRemainder);
                var workerBundleStart = context.SourceStartBundleIndex + localWorkerBundleStart;
                var workerBundleCount = workerIndex < context.CopyBundlesPerWorkerRemainder ? context.CopyBundlesPerWorker + 1 : context.CopyBundlesPerWorker;
                var workerConstraintStart = workerBundleStart << BundleIndexing.VectorShift;
                // 请注意,我们可以迭代的约束数量受类型Batch的约束计数的限制。最后一个捆绑包可能未满。
                var workerConstraintCount = Math.Min(context.TypeBatchConstraintCount - workerConstraintStart, workerBundleCount << BundleIndexing.VectorShift);
                if (workerConstraintCount <= 0)
                    return; // 没有剩余的工作。
                var localWorkerConstraintStart = localWorkerBundleStart << BundleIndexing.VectorShift;

                ref var typeBatch = ref *typeBatchPointer;
                solver.TypeProcessors[typeBatch.TypeId].CopyToCache(ref typeBatch,
                    workerBundleStart, localWorkerBundleStart, workerBundleCount,
                    workerConstraintStart, localWorkerConstraintStart, workerConstraintCount,
                    ref context.IndexToHandleCache, ref context.PrestepDataCache, ref context.AccumulatesImpulsesCache);
            }
        }

        unsafe void CopyToCacheAndSort(BufferPool pool)
        {
            LSBRadixSort.Sort(
                ref context.SortKeys, ref context.SourceIndices,
                ref context.ScratchKeys, ref context.ScratchValues, 0, context.ConstraintsInSortRegionCount,
                context.KeyUpperBound, pool,
                out context.SortedKeys, out context.SortedSourceIndices);

            var workerBundleStart = context.SourceStartBundleIndex;
            var workerBundleCount = 0 < context.CopyBundlesPerWorkerRemainder ? context.CopyBundlesPerWorker + 1 : context.CopyBundlesPerWorker;
            var workerConstraintStart = workerBundleStart << BundleIndexing.VectorShift;
            // 请注意,我们可以迭代的约束数量受类型Batch的约束计数的限制。最后一个捆绑包可能未满。
            var workerConstraintCount = Math.Min(context.TypeBatchConstraintCount - workerConstraintStart, workerBundleCount << BundleIndexing.VectorShift);
            if (workerConstraintCount <= 0)
                return; // 没有剩余的工作。

            ref var typeBatch = ref *typeBatchPointer;
            solver.TypeProcessors[typeBatch.TypeId].CopyToCache(ref typeBatch,
                workerBundleStart, 0, workerBundleCount,
                workerConstraintStart, 0, workerConstraintCount,
                ref context.IndexToHandleCache, ref context.PrestepDataCache, ref context.AccumulatesImpulsesCache);
        }

        unsafe void Regather(int workerIndex)
        {
            var localWorkerBundleStart = context.BundlesPerWorker * workerIndex + Math.Min(workerIndex, context.BundlesPerWorkerRemainder);
            var workerBundleStart = context.SourceStartBundleIndex + localWorkerBundleStart;
            var workerBundleCount = workerIndex < context.BundlesPerWorkerRemainder ? context.BundlesPerWorker + 1 : context.BundlesPerWorker;
            var workerConstraintStart = workerBundleStart << BundleIndexing.VectorShift;
            // 请注意,我们可以迭代的约束数量受类型Batch的约束计数的限制。最后一个捆绑包可能未满。
            var workerConstraintCount = Math.Min(context.TypeBatchConstraintCount - workerConstraintStart, workerBundleCount << BundleIndexing.VectorShift);
            if (workerConstraintCount <= 0)
                return; // 没有剩余的工作。
            var localWorkerConstraintStart = localWorkerBundleStart << BundleIndexing.VectorShift;
            ref var firstSourceIndex = ref context.SortedSourceIndices[localWorkerConstraintStart];

            ref var typeBatch = ref *typeBatchPointer;
            solver.TypeProcessors[typeBatch.TypeId].Regather(ref typeBatch, workerConstraintStart, workerConstraintCount, ref firstSourceIndex,
                ref context.IndexToHandleCache, ref context.BodyReferencesCache, ref context.PrestepDataCache, ref context.AccumulatesImpulsesCache, ref handlesToConstraints);

        }



        unsafe void SortByBodyLocation(TypeBatch* typeBatch, int bundleStartIndex, int constraintCount, Buffer<ConstraintLocation> handlesToConstraints, int bodyCount,
            BufferPool pool, IThreadDispatcher threadDispatcher)
        {
            int bundleCount = (constraintCount >> BundleIndexing.VectorShift);
            if ((constraintCount & BundleIndexing.VectorMask) != 0)
                ++bundleCount;
            var threadCount = threadDispatcher == null ? 1 : threadDispatcher.ThreadCount;

            pool.TakeAtLeast(constraintCount, out context.SourceIndices);
            pool.TakeAtLeast(constraintCount, out context.SortKeys);
            pool.TakeAtLeast(constraintCount, out context.ScratchKeys);
            pool.TakeAtLeast(constraintCount, out context.ScratchValues);
            pool.TakeAtLeast(constraintCount, out context.IndexToHandleCache);

            var typeProcessor = solver.TypeProcessors[typeBatch->TypeId];
            typeProcessor.GetBundleTypeSizes(out var bodyReferencesBundleSize, out var prestepBundleSize, out var accumulatedImpulseBundleSize);

            // Worker调用的TypeBatch会将正文引用强制转换为适当的类型。
            // 使用无类型缓冲区可以轻松地将缓冲区缓存在约束优化器中,而不是缓存在单个类型批中。
            pool.TakeAtLeast(bundleCount * bodyReferencesBundleSize, out context.BodyReferencesCache);
            pool.TakeAtLeast(bundleCount * prestepBundleSize, out context.PrestepDataCache);
            pool.TakeAtLeast(bundleCount * accumulatedImpulseBundleSize, out context.AccumulatesImpulsesCache);

            context.BundlesPerWorker = bundleCount / threadCount;
            context.BundlesPerWorkerRemainder = bundleCount - context.BundlesPerWorker * threadCount;
            context.TypeBatchConstraintCount = typeBatch->ConstraintCount;
            context.SourceStartBundleIndex = bundleStartIndex;

            // 第二阶段使用一个工作器进行排序。
            if (threadCount > 1)
            {
                context.CopyBundlesPerWorker = bundleCount / (threadCount - 1);
                context.CopyBundlesPerWorkerRemainder = bundleCount - context.CopyBundlesPerWorker * (threadCount - 1);
            }
            else
            {
                // 如果只有一个Worker(单线程运行时就是这种情况),Worker将不得不进行排序和复制。
                context.CopyBundlesPerWorker = bundleCount;
                context.CopyBundlesPerWorkerRemainder = 0;
            }
            context.ConstraintsInSortRegionCount = constraintCount;
            context.KeyUpperBound = bodyCount - 1;

            this.typeBatchPointer = typeBatch;
            this.threadDispatcher = threadDispatcher;
            this.handlesToConstraints = handlesToConstraints;

            if (threadDispatcher == null)
            {
                GenerateSortKeys(0);
                CopyToCacheAndSort(pool);
                Regather(0);
            }
            else
            {
                threadDispatcher.DispatchWorkers(generateSortKeysDelegate);
                threadDispatcher.DispatchWorkers(copyToCacheAndSortDelegate);
                threadDispatcher.DispatchWorkers(regatherDelegate);
            }

            this.typeBatchPointer = null;
            this.threadDispatcher = null;
            this.handlesToConstraints = new Buffer<ConstraintLocation>();

            // 这是一个纯调试函数。
            solver.TypeProcessors[typeBatch->TypeId].VerifySortRegion(ref *typeBatch, bundleStartIndex, constraintCount, ref context.SortedKeys, ref context.SortedSourceIndices);

            pool.Return(ref context.SourceIndices);
            pool.Return(ref context.SortKeys);
            pool.Return(ref context.ScratchKeys);
            pool.Return(ref context.ScratchValues);
            pool.Return(ref context.IndexToHandleCache);
            pool.Return(ref context.BodyReferencesCache);
            pool.Return(ref context.PrestepDataCache);
            pool.Return(ref context.AccumulatesImpulsesCache);
        }
    }
}
